# -*- coding: utf-8 -*-

import argparse
from deep_sort_wrapper import DeepSORT
import cv2
import os
import sys
import time

this_file_path = os.path.abspath(__file__)
this_file_dir, filename = os.path.split(this_file_path)
yolo_path = this_file_dir + "/keras_yolo3"
sys.path.append(yolo_path)

from yolo import YOLO


def main():
    args = parse_args()

    yolo = YOLO(**vars(args))
    deep_sort = DeepSORT(yolo, **vars(args))

    if not ("input" in args):
        print("Must specify at least video_input_path.  See usage with --help.")
        exit(1)

    vid = cv2.VideoCapture(args.input)
    if not vid.isOpened():
        print("Couldn't open webcam or video")
        exit(1)

    video_FourCC = int(vid.get(cv2.CAP_PROP_FOURCC))
    video_fps = vid.get(cv2.CAP_PROP_FPS)
    video_size = (int(vid.get(cv2.CAP_PROP_FRAME_WIDTH)),
                  int(vid.get(cv2.CAP_PROP_FRAME_HEIGHT)))
    isOutput = True if args.output != "" else False
    if isOutput:
        print("!!! TYPE:", type(args.output), type(video_FourCC), type(video_fps), type(video_size))
        out = cv2.VideoWriter(args.output, video_FourCC, video_fps, video_size)

    total_time = 0
    frame_count = 0
    while True:
        start_time = time.time()
        return_value, frame = vid.read()
        if not return_value:
            cv2.waitKey(0)
            break

        tracks = deep_sort.process_frame(frame)

        for bbox in deep_sort.bboxes:
            draw_bbox(frame, bbox, (255, 0, 0), 2)

        for track in tracks:
            bbox = track[2:6]
            draw_bbox(frame, bbox, (255, 255, 255), 2)
            cv2.putText(frame, str(track[1]), (bbox[0], bbox[1]), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 255, 0), 2)

        cv2.imshow("result", frame)
        if isOutput:
            out.write(frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

        finish_time = time.time()
        proc_time = finish_time - start_time
        frame_count += 1
        if (frame_count > 1):
            # skip first frame, because it take too much time then other
            total_time += proc_time
        fps = 1/proc_time
        print("fps: %f" % (fps))


    frame_count -= 1 # remove first frame
    avg_fps = frame_count / total_time
    print("average fps: %f" % (avg_fps))
    yolo.close_session()


def draw_bbox(img, bbox, color, thickness):
    x1, y1 = bbox[:2]
    w, h = bbox[2:]
    x2 = x1 + w
    y2 = y1 + h
    cv2.rectangle(img, (x1, y1), (x2, y2), color, thickness)


def bool_string(input_string):
    if input_string not in {"True", "False"}:
        raise ValueError("Please Enter a valid Ture/False choice")
    else:
        return (input_string == "True")


def parse_args():
    """
    Parse command line arguments.
    """
    parser = argparse.ArgumentParser(description="Deep SORT with YOLOv3")

    '''
    Command line options for YOLO
    '''
    parser.add_argument(
        '--model', type=str, default='yolo.h5',
        help='path to YOLO model weight file'
    )

    parser.add_argument(
        '--anchors', type=str, default='keras_yolo3/model_data/yolo_anchors.txt',
        help='path to YOLO anchor definitions'
    )

    parser.add_argument(
        '--classes', type=str, default='keras_yolo3/model_data/coco_classes.txt',
        help='path to YOLO class definitions'
    )

    parser.add_argument(
        '--tracking_classes', type=str, default='tracking_classes.txt',
        help='path to class for tracking'
    )

    parser.add_argument(
        '--gpu_num', type=int, default=1,
        help='Number of GPU to use for YOLO'
    )

    parser.add_argument(
        '--model_image_size', type=int, default=None,
        help='image size used by YOLO detection, Multiples of 32 require'
    )

    parser.add_argument(
        '--score', type=float, default=0.3,
        help='YOLO detection score_threshold'
    )

    parser.add_argument(
        '--iou', type=float, default=0.45,
        help='YOLO detection iou_threshold'
    )

    '''
    Command line options for DeepSORT
    '''
    parser.add_argument(
        "--feature_model_filename", type=str, default="resources/networks/mars-small128.pb",
        help="Model filename for feature detector."
    )
    # parser.add_argument(
    #     "--nms_max_overlap", type=float, default=1.0,
    #     help=""
    # )
    parser.add_argument(
        "--min_detection_height", default=0, type=int,
        help="Threshold on the detection bounding box height. "
             "Detections with height smaller than this value are disregarded"
    )
    parser.add_argument(
        "--max_cosine_distance", type=float, default=0.2,
        help="Gating threshold for cosine distance metric (object appearance)."
    )
    parser.add_argument(
        "--nn_budget", type=int, default=None,
        help="Maximum size of the appearance descriptors gallery. "
             "If None, no budget is enforced."
    )
    parser.add_argument(
        "--update_step", type=int, default=1,
        help="Update step for tracker."
    )

    '''
    Command line options for main program
    '''
    parser.add_argument(
        "--input", type=str, required=True,
        help="Video input path, OpenCV File Sequence Format also can be used."
    )

    parser.add_argument(
        "--output", type=str, default="",
        help="[Optional] Video output path"
    )

    # parser.add_argument(
    #     "--output_file", type=str, default="",
    #     help="Path to the tracking output file. This file will"
    #          " contain the tracking results on completion."
    # )

    # parser.add_argument(
    #     "--display", default=True, type=bool_string,
    #     help="Show intermediate tracking results"
    # )

    FLAGS = parser.parse_args()
    FLAGS.model_image_size = (FLAGS.model_image_size, FLAGS.model_image_size)
    print(FLAGS)
    return FLAGS


if __name__ == "__main__":
    main()
